#include <iostream>
#include <chrono>
#include <random>
#include <Eigen/Geometry>
#include <omp.h>

#include "render.hpp"
#include "asset.hpp"
#include "output.hpp"

static const char * getenv_or_panic(const char * name)
{
    const char * env = getenv(name);
    if (!env)
    {
        std::cerr << "Environment variable `" << name << "` cannot be empty\n";
        exit(1);
    }
    return env;
}

static const char * getenv_or(const char * name, const char * def)
{
    const char * env = getenv(name);
    return env ? env : def;
}

class ProgressBar
{
    int a_total;
    int a_count;
    int d_total;
    int d_count;
    void show() const
    {
        std::cerr << "|";
        for (int i = 0; i < d_count; i++)
            std::cerr.put('*');
        for (int i = d_count; i < d_total; i++)
            std::cerr.put(' ');
        std::cerr << "|\r";
    }
public:
    ProgressBar(int total, int width = 50):
        a_total(total), a_count(0),
        d_total(width), d_count(0)
    {
        show();
    }
    void advance()
    {
        a_count++;
        int n_d_count = a_count * d_total / a_total;
        if (n_d_count > d_count)
        {
            d_count = n_d_count;
            show();
        }
    }
};

int main()
{
    std::vector<Object> objects(3);

    // object

    objects[0].type = 1;
    objects[0].brdf = [](
        const Eigen::Vector3f & color,
        const Eigen::Vector3f & in_dir,
        const Eigen::Vector3f & out_dir,
        const Eigen::Vector3f & normal)
    {
        // Eigen::Vector3f half_dir = (in_dir + out_dir).normalized();
        float diff = 1.0f / (float)EIGEN_PI;
        // float spec = powf(fmaxf(normal.dot(half_dir), 0.0f), 64.0f) / (float)EIGEN_PI;
        // return color * (diff * 0.7f + spec * 0.3f);
        return color * diff;
    };

    load_mesh(getenv_or_panic("APP_MESH_PATH"), objects[0].vertices, objects[0].faces);

    Eigen::RowVector3f min = objects[0].vertices.colwise().minCoeff();
    Eigen::RowVector3f max = objects[0].vertices.colwise().maxCoeff();
    Eigen::RowVector3f center = (min + max) / 2.0f;
    float size = (max - min).maxCoeff() / 2.0f;
    objects[0].vertices.rowwise() -= center;
    objects[0].vertices /= size;

    const char * rotation_str = getenv("APP_MESH_ROTATION");
    if (rotation_str)
    {
        int order[3];
        float degree[3];
        sscanf(rotation_str, "%d,%d,%d:%f,%f,%f",
            &order[0], &order[1], &order[2], &degree[0], &degree[1], &degree[2]);
        for (int i = 0; i < 3; i++)
        {
            Eigen::AngleAxisf rot(degree[i] / 180.0f * (float)EIGEN_PI, Eigen::Vector3f::Unit(order[i]));
            objects[0].vertices *= rot.toRotationMatrix().transpose();
        }
    }

    std::cerr << "#V = " << objects[0].vertices.rows() <<
        ", #F = " << objects[0].faces.rows() << "\n";

    // light

    objects[1].type = 2;
    objects[1].light_color = Eigen::Vector3f::Ones() * 50.0;
    // objects[1].light_color = Eigen::Vector3f(0.0f, 1.0f, 1.0f) * 50.0;
    objects[1].vertices.resize(4, 3);
    objects[1].vertices <<
        -0.5f, 3.0f, 4.0f,
        -0.5f, 4.0f, 4.0f,
         0.5f, 3.0f, 4.0f,
         0.5f, 4.0f, 4.0f;
    objects[1].faces.resize(2, 3);
    objects[1].faces <<
        0, 1, 2,
        2, 1, 3;

    // object 2

    objects[2].type = 1;
    objects[2].brdf = objects[0].brdf;
    objects[2].vertices.resize(4, 3);
    objects[2].vertices <<
        -5.0f, -5.0f, -2.0f,
        -5.0f,  5.0f, -2.0f,
         5.0f, -5.0f, -2.0f,
         5.0f,  5.0f, -2.0f;
    objects[2].faces.resize(2, 3);
    objects[2].faces <<
        2, 1, 0,
        3, 1, 2;

    // build scene

    auto begin_time = std::chrono::steady_clock::now();

    Scene scene(std::move(objects), [](const Eigen::Vector3f &)
    {
        return Eigen::Vector3f::Zero();
        // return Eigen::Vector3f::Ones() * 0.5f;
    });

    auto end_time = std::chrono::steady_clock::now();
    std::chrono::duration<double> duration = end_time - begin_time;
    std::cerr << "build duration = " << duration.count() << "s\n";

    // camera

    Eigen::Vector3f cam_pos = Eigen::Vector3f(0.0f, 0.0f, 3.0f);
    Eigen::Vector3f cam_dir = Eigen::Vector3f(0.0f, 0.0f, -1.0f);
    Eigen::Vector3f cam_up = Eigen::Vector3f(0.0f, 1.0f, 0.0f);

    const float TW = 1.0f, TH = 1.0f;
    const int W = 300, H = 300;
    const float PW = TW / W, PH = TH / H;
    const float TX = -TW / 2.0f, TY = -TH / 2.0f;
    const float D = 1.0f;
    static uint8_t img[H][W][3];

    // render

    begin_time = std::chrono::steady_clock::now();

    std::mt19937 rng(atoi(getenv_or("APP_RNG_SEED", "1")));
    std::vector<int> seeds(H);
    for (int & seed : seeds)
        seed = rng();

    const char * num_threads_str = getenv("APP_OMP_NUM_THREADS");
    if (num_threads_str)
        omp_set_num_threads(atoi(num_threads_str));

    ProgressBar bar(H * W);

#pragma omp parallel for schedule(dynamic)
    for (int y = 0; y < H; y++)
    {
        std::mt19937 row_rng(seeds[y]);
        std::uniform_real_distribution<float> dis;

        RenderConfig config =
        {
            atoi(getenv_or_panic("APP_RENDER_SPP")),
            (float)atof(getenv_or_panic("APP_RENDER_BOUNCE")),
            [&dis, &row_rng](){ return dis(row_rng); }
        };

        for (int x = 0; x < W; x++)
        {
            Eigen::Vector3f color = shade_pixel(scene, cam_pos, cam_dir, cam_up, D,
                TX + x * PW, TY + y * PH, PW, PH, config) * 255.0f;
            color = color.cwiseMax(0.0f);
            color = color.cwiseMin(255.0f);
            for (int i = 0; i < 3; i++)
                img[H - 1 - y][x][i] = color(i);

#pragma omp critical
                bar.advance();
        }
    }

    std::cerr << "\n";

    end_time = std::chrono::steady_clock::now();
    duration = end_time - begin_time;
    std::cerr << "render duration = " << duration.count() << "s\n";

    // save image

    output_png(getenv_or("APP_OUTPUT_PATH", "output.png"), W, H, img);

    return 0;
}
